from __future__ import print_function
import sys
import torch
from global_var import BACKEND, DATA_PATH, TimeUse
if BACKEND == "npu":
    import torch_npu
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import random
import os
import sys
import argparse
import numpy as np
from InceptionResNetV2_pl import *
from sklearn.mixture import GaussianMixture
import dataloader_webvision_pl as dataloader
import torchnet
import copy
import functools
from apex import amp

parser = argparse.ArgumentParser(description='PyTorch WebVision Training')
parser.add_argument('--batch_size', default=32, type=int, help='train batchsize')
parser.add_argument('--lr', '--learning_rate', default=0.01, type=float, help='initial learning rate')
parser.add_argument('--alpha', default=0.5, type=float, help='parameter for Beta')
parser.add_argument('--lambda_u', default=0, type=float, help='weight for unsupervised loss')
parser.add_argument('--p_threshold', default=0.5, type=float, help='clean probability threshold')
parser.add_argument('--T', default=0.5, type=float, help='sharpening temperature')
parser.add_argument('--num_epochs', default=110, type=int)
parser.add_argument('--id', default='',type=str)
parser.add_argument('--seed', default=0)
parser.add_argument('--gpuid1', default=0, type=int)
parser.add_argument('--gpuid2', default=1, type=int)
parser.add_argument('--amp', default=True, type=bool)
parser.add_argument('--num_class', default=50, type=int)
parser.add_argument('--data_path', default='../../data/webvision/', type=str, help='path to dataset')
# bpl
parser.add_argument('--is_load', default=True, type=bool)
parser.add_argument('--load_epoch', default=94, type=int)
parser.add_argument('--save_epoch', default=200, type=int)
parser.add_argument('--save_interval', default=5, type=int)
parser.add_argument('--pl_start_epoch', default=1, type=int)
parser.add_argument('--pl_end_epoch', default=110, type=int)
args = parser.parse_args()
args.data_path = DATA_PATH


if BACKEND == "npu":
    os.environ["ASCEND_RT_VISIBLE_DEVICES"] = '%s,%s'%(args.gpuid1, args.gpuid2)
    device1 = torch.device("npu:0")
    device2 = torch.device("npu:1")
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.npu.manual_seed_all(args.seed)
elif BACKEND == "cuda":
    os.environ["CUDA_VISIBLE_DEVICES"] = '%s,%s'%(args.gpuid1, args.gpuid2)
    device1 = torch.device("cuda:0")
    device2 = torch.device("cuda:1")
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

def nn_mixup_two_targets(x1, x2, y1, y2, alpha=1.0, device='cuda', is_bias=False):
    """Returns nn-mixed inputs, pairs of targets, and lambda
    """
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1
    if is_bias: lam = max(lam, 1-lam)

    mixed_x = lam*x1 + (1-lam)*x2
    y_a, y_b = y1, y2 
    return mixed_x, y_a, y_b, lam

def mixup_ce_loss_with_softmax(preds, targets_a, targets_b, lam, mask=None):
    """ mixed categorical cross-entropy loss
    """
    if mask is not None:
        mixup_loss_a = -torch.mean(mask * torch.sum(F.softmax(targets_a,1) * F.log_softmax(preds, dim=1), dim=1))
        mixup_loss_b = -torch.mean(mask * torch.sum(F.softmax(targets_b,1) * F.log_softmax(preds, dim=1), dim=1))
    else:
        mixup_loss_a = -torch.mean(torch.sum(F.softmax(targets_a, 1) * F.log_softmax(preds, dim=1), dim=1))
        mixup_loss_b = -torch.mean(torch.sum(F.softmax(targets_b, 1) * F.log_softmax(preds, dim=1), dim=1))

    mixup_loss = lam* mixup_loss_a + (1 - lam) * mixup_loss_b
    return mixup_loss

# Training
def train(epoch,net,net2,optimizer,labeled_trainloader,unlabeled_trainloader):
    net.train()
    net2.eval() #fix one network and train the other
    
    unlabeled_train_iter = iter(unlabeled_trainloader)    
    num_iter = (len(labeled_trainloader.dataset)//args.batch_size)+1
    for batch_idx, (inputs_x, inputs_x2, labels_x, w_x) in enumerate(labeled_trainloader):      
        try:
            inputs_u, inputs_u2 = unlabeled_train_iter.next()
        except:
            unlabeled_train_iter = iter(unlabeled_trainloader)
            inputs_u, inputs_u2 = unlabeled_train_iter.next()                 
        batch_size = inputs_x.size(0)
        
        # Transform label to one-hot
        labels_x = torch.zeros(batch_size, args.num_class).scatter_(1, labels_x.view(-1,1), 1)        
        w_x = w_x.view(-1,1).type(torch.FloatTensor) 

        inputs_x, inputs_x2, labels_x, w_x = inputs_x.to(device), inputs_x2.to(device), labels_x.to(device), w_x.to(device)
        inputs_u, inputs_u2 = inputs_u.to(device), inputs_u2.to(device)

        with torch.no_grad():
            # label co-guessing of unlabeled samples
            outputs_u11 = net(inputs_u)
            outputs_u12 = net(inputs_u2)
            outputs_u21 = net2(inputs_u)
            outputs_u22 = net2(inputs_u2)            
            
            pu = (torch.softmax(outputs_u11, dim=1) + torch.softmax(outputs_u12, dim=1) + torch.softmax(outputs_u21, dim=1) + torch.softmax(outputs_u22, dim=1)) / 4       
            ptu = pu**(1/args.T) # temparature sharpening
            
            targets_u = ptu / ptu.sum(dim=1, keepdim=True) # normalize
            targets_u = targets_u.detach()       
            
            # label refinement of labeled samples
            outputs_x = net(inputs_x)
            outputs_x2 = net(inputs_x2)            
            
            px = (torch.softmax(outputs_x, dim=1) + torch.softmax(outputs_x2, dim=1)) / 2
            px = w_x*labels_x + (1-w_x)*px              
            ptx = px**(1/args.T) # temparature sharpening 
                       
            targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize           
            targets_x = targets_x.detach()       
        
        # mixmatch
        l = np.random.beta(args.alpha, args.alpha)        
        l = max(l, 1-l)
                
        all_inputs = torch.cat([inputs_x, inputs_x2, inputs_u, inputs_u2], dim=0)
        all_targets = torch.cat([targets_x, targets_x, targets_u, targets_u], dim=0)

        idx = torch.randperm(all_inputs.size(0))

        input_a, input_b = all_inputs, all_inputs[idx]
        target_a, target_b = all_targets, all_targets[idx]

        mixed_input = l * input_a[:batch_size*2] + (1 - l) * input_b[:batch_size*2]        
        mixed_target = l * target_a[:batch_size*2] + (1 - l) * target_b[:batch_size*2]
                
        logits = net(mixed_input)
        
        Lx = -torch.mean(torch.sum(F.log_softmax(logits, dim=1) * mixed_target, dim=1))
        
        prior = torch.ones(args.num_class)/args.num_class
        prior = prior.to(device)        
        pred_mean = torch.softmax(logits, dim=1).mean(0)
        penalty = torch.sum(prior*torch.log(prior/pred_mean))
       
        loss = Lx + penalty
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        sys.stdout.write('\r')
        sys.stdout.write('%s | Epoch [%3d/%3d] Iter[%4d/%4d]\t Labeled loss: %.2f'
                %(args.id, epoch, args.num_epochs, batch_idx+1, num_iter, Lx.item()))
        sys.stdout.flush()

# Training with BPL
def train_bpl(epoch, net, net2, optimizer, bpl_trainloader, device, whichnet):
    criterion = SemiLoss()   
    
    net.train()
    net2.eval() #fix one network and train the other
    print("{} Train Bpl Epoch {}".format(whichnet, epoch))

    is_bpl = True if epoch > args.pl_start_epoch and epoch <= args.pl_end_epoch else False
    if is_bpl: 
        # + copy model
        net_pl = copy.deepcopy(net)
        net_pl = net_pl.train()
        optimizer_pl = torch.optim.SGD(net_pl.parameters(), lr=0.02)
        optimizer_pl.load_state_dict(optimizer.state_dict())  
        net_pl, optimizer_pl = amp.initialize(net_pl, optimizer_pl, opt_level="O1")
    for batch_idx, (inputs_x, inputs_x2, labels_x, w_x, 
                    inputs_u, inputs_u2,
                    inputs_nn_x, inputs_nn_u,
                    mask_x, mask_u,
                    umask_x, umask_u) in enumerate(bpl_trainloader):
        # print("Step {}".format(batch_idx))
        batch_size = inputs_x.size(0)
        # Transform label to one-hot
        labels_x = torch.zeros(batch_size, args.num_class).scatter_(1, labels_x.view(-1,1), 1)        
        w_x = w_x.view(-1,1).type(torch.FloatTensor)

        inputs_x, inputs_x2, labels_x, w_x = inputs_x.to(device,non_blocking=True), inputs_x2.to(device,non_blocking=True), labels_x.to(device,non_blocking=True), w_x.to(device,non_blocking=True)
        inputs_u, inputs_u2 = inputs_u.to(device, non_blocking=True), inputs_u2.to(device, non_blocking=True)
        inputs_nn_x, inputs_nn_u = inputs_nn_x.to(device, non_blocking=True), inputs_nn_u.to(device, non_blocking=True)
        mask_x, mask_u = mask_x.to(device, non_blocking=True), mask_u.to(device, non_blocking=True)
        umask_x, umask_u = umask_x.to(device, non_blocking=True), umask_u.to(device, non_blocking=True)
        umask_x, umask_u = umask_x.view((batch_size, -1)), umask_u.view((batch_size, -1))
        # print("label co-guessing start.")
        with torch.no_grad():
            # label co-guessing of unlabeled samples
            outputs_u11 = net(inputs_u)
            outputs_u12 = net(inputs_u2)
            outputs_u21 = net2(inputs_u)
            outputs_u22 = net2(inputs_u2)            
            
            pu = (torch.softmax(outputs_u11, dim=1) + torch.softmax(outputs_u12, dim=1) + torch.softmax(outputs_u21, dim=1) + torch.softmax(outputs_u22, dim=1)) / 4       
            ptu = pu**(1/args.T) # temparature sharpening
            
            targets_u = ptu / ptu.sum(dim=1, keepdim=True) # normalize
            targets_u = targets_u.detach()       
            
            # label refinement of labeled samples
            outputs_x = net(inputs_x)
            outputs_x2 = net(inputs_x2)            
            
            px = (torch.softmax(outputs_x, dim=1) + torch.softmax(outputs_x2, dim=1)) / 2
            px = w_x*labels_x + (1-w_x)*px              
            ptx = px**(1/args.T) # temparature sharpening 
                       
            targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize           
            targets_x = targets_x.detach()
        # print("label co-guessing ends.")

        # mixmatch
        l = np.random.beta(args.alpha, args.alpha)        
        l = max(l, 1-l)
                
        all_inputs = torch.cat([inputs_x, inputs_x2, inputs_u, inputs_u2], dim=0)
        all_targets = torch.cat([targets_x, targets_x, targets_u, targets_u], dim=0)

        idx = torch.randperm(all_inputs.size(0))

        input_a, input_b = all_inputs, all_inputs[idx]
        target_a, target_b = all_targets, all_targets[idx]

        mixed_input = l * input_a[:batch_size*2] + (1 - l) * input_b[:batch_size*2]        
        mixed_target = l * target_a[:batch_size*2] + (1 - l) * target_b[:batch_size*2]

        if is_bpl:
            # + bpl model training
            # ++ compute supervised loss
            logits_pl = net_pl(mixed_input)
        
            Lx_pl = -torch.mean(torch.sum(F.log_softmax(logits_pl, dim=1) * mixed_target, dim=1))
            
            prior = torch.ones(args.num_class)/args.num_class
            prior = prior.to(device)        
            pred_mean = torch.softmax(logits_pl, dim=1).mean(0)
            penalty_pl = torch.sum(prior*torch.log(prior/pred_mean))
        
            loss_pl_sup = Lx_pl + penalty_pl
            loss_pl = loss_pl_sup
            # -- compute supervised loss

            # ++ compute bpl regularization loss
            model_ema2 = net2
            # 随机选batch_size大小
            data_idx = torch.randperm(batch_size * 2)[:batch_size]
            data_unreli = torch.cat([inputs_nn_x, inputs_nn_u], dim=0)[data_idx]
            data_reli = torch.cat([inputs_x, inputs_u], dim=0)[data_idx]
            data_mask = torch.cat([mask_x, mask_u])[data_idx]
            data_reli_list = [data_reli, ]
            # +++ mean-teacher
            with torch.no_grad():
                ema_outputs_unreli = model_ema2(data_unreli)
                ema_outputs_reli_list = list(map(lambda x: model_ema2(x), data_reli_list))
            # --- mean-teacher

            # +++ mix up for consistency loss
            partial_nn_mixup_two_targets = functools.partial(nn_mixup_two_targets, x2=data_unreli, y2=ema_outputs_unreli, alpha=1.0, device=device, is_bias=False)
            results_nn_mixup_two_targets = list(map(lambda x, y: partial_nn_mixup_two_targets(x1=x, y1=y), data_reli_list, ema_outputs_reli_list))
            mixed_outputs_u_list = list(map(lambda x: net_pl(x[0]), results_nn_mixup_two_targets))
            mix_loss_list = list(map(lambda x, y: mixup_ce_loss_with_softmax(x, y[1], y[2], y[3], data_mask),
                                     mixed_outputs_u_list, results_nn_mixup_two_targets))
            loss_pl_reg = sum(mix_loss_list) / len(mix_loss_list)
            # --- mix up for consistency loss

            # -- compute bpl regularization loss

            loss_pl = loss_pl_sup + loss_pl_reg

            optimizer_pl.zero_grad()
            if args.amp:
                with amp.scale_loss(loss_pl, optimizer_pl) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss_pl.backward()
            optimizer_pl.step()
            # - bpl model training      

            # print("label co-guessing start.")
            with torch.no_grad():
                # label co-guessing of unlabeled samples
                outputs_u11 = (1 - umask_u) * outputs_u11 + umask_u * net_pl(inputs_u)
                outputs_u12 = (1 - umask_u) * outputs_u12 + umask_u * net_pl(inputs_u2)
                outputs_u21 = net2(inputs_u)
                outputs_u22 = net2(inputs_u2)

                pu = (torch.softmax(outputs_u11, dim=1) + torch.softmax(outputs_u12, dim=1) + \
                      torch.softmax(outputs_u21, dim=1) + torch.softmax(outputs_u22, dim=1)) / 4
                ptu = pu ** (1 / args.T)  # temparature sharpening

                targets_u = ptu / ptu.sum(dim=1, keepdim=True)  # normalize
                targets_u = targets_u.detach()

                # label refinement of labeled samples
                outputs_x = (1 - umask_x) * outputs_x + umask_x * net_pl(inputs_x)
                outputs_x2 = (1 - umask_x) * outputs_x2 + umask_x * net_pl(inputs_x2)
                px = (torch.softmax(outputs_x, dim=1) + torch.softmax(outputs_x2, dim=1)) / 2
                px = w_x * labels_x + (1 - w_x) * px
                ptx = px ** (1 / args.T)  # temparature sharpening
                targets_x = ptx / ptx.sum(dim=1, keepdim=True)  # normalize
                targets_x = targets_x.detach()
            # print("label co-guessing ends.")
            all_targets = torch.cat([targets_x, targets_x, targets_u, targets_u], dim=0)
            target_a, target_b = all_targets, all_targets[idx]
            mixed_target = l * target_a[:batch_size * 2] + (1 - l) * target_b[:batch_size * 2]

        logits = net(mixed_input)
        
        Lx = - torch.mean(torch.sum(F.log_softmax(logits, dim=1) * mixed_target, dim=1))
        
        prior = torch.ones(args.num_class)/args.num_class
        prior = prior.to(device)        
        pred_mean = torch.softmax(logits, dim=1).mean(0)
        penalty = torch.sum(prior*torch.log(prior/pred_mean))
       
        loss = Lx + penalty
        # compute gradient and do SGD step
        optimizer.zero_grad()
        if args.amp:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        optimizer.step()
        
        # sys.stdout.write('\n')
        # sys.stdout.write('%s |%s Epoch [%3d/%3d] Iter[%4d/%4d]\t Labeled loss: %.2f'
        #         %(args.id, whichnet, epoch, args.num_epochs, batch_idx+1, len(bpl_trainloader), Lx.item()))
        # sys.stdout.flush()

def warmup(epoch,net,optimizer,dataloader,device):
    net.train()
    num_iter = (len(dataloader.dataset)//dataloader.batch_size)+1
    for batch_idx, (inputs, labels, path) in enumerate(dataloader):      
        inputs, labels = inputs.to(device), labels.to(device) 
        optimizer.zero_grad()
        outputs = net(inputs)               
        loss = CEloss(outputs, labels)   
        
        #penalty = conf_penalty(outputs)
        L = loss #+ penalty      

        L.backward()  
        optimizer.step() 

        # sys.stdout.write('\r')
        # sys.stdout.write('%s | Epoch [%3d/%3d] Iter[%4d/%4d]\t CE-loss: %.4f'
        #         %(args.id, epoch, args.num_epochs, batch_idx+1, num_iter, loss.item()))
        # sys.stdout.flush()
        
        
def test(epoch, net1, net2, device1, device2, test_loader):
    acc_meter.reset()
    net1.eval()
    net2.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            inputs_1 = inputs.to(device1)
            inputs_2 = inputs.to(device2)
            outputs1 = net1(inputs_1)
            outputs2 = net2(inputs_2)
            outputs = outputs1.cpu() + outputs2.cpu()
            _, predicted = torch.max(outputs, 1)                 
            acc_meter.add(outputs, targets)
    accs = acc_meter.value()
    return accs


def eval_train(eval_loader,model,device,whichnet):   
    CE = nn.CrossEntropyLoss(reduction='none')
    model.eval()
    num_iter = (len(eval_loader.dataset)//eval_loader.batch_size)+1
    losses = torch.zeros(len(eval_loader.dataset))   
    uncertainty = torch.zeros(len(eval_loader.dataset))
    embeddings = torch.zeros((len(eval_loader.dataset), 1536)) 
    with torch.no_grad():
        for batch_idx, (inputs, targets, index) in enumerate(eval_loader):
            inputs, targets = inputs.to(device), targets.to(device,non_blocking=True) 
            outputs = model(inputs)
            loss = CE(outputs, targets)
            if BACKEND == "npu":
                batch_uncertainty = - (F.softmax(outputs, dim=1) * torch.log(F.softmax(outputs, dim=1))).sum(dim=1)
            else:
                batch_uncertainty = - (F.softmax(outputs, dim=1) * torch.log_softmax(outputs, dim=1)).sum(dim=1)

            batch_embeddings = torch.flatten(
                model.last_k_layer_forward(
                    inputs, k=-2, lin=0, lout=5
                ), 
                start_dim=1, end_dim=-1)

            losses[index] = loss.clone().detach().cpu()
            uncertainty[index] = batch_uncertainty.clone().detach().cpu()
            embeddings[index, :] = batch_embeddings.clone().detach().cpu()     
            # sys.stdout.write('\n')
            # sys.stdout.write('|%s Evaluating loss Iter[%3d/%3d]\t' %(whichnet,batch_idx,num_iter))
            # sys.stdout.flush()

    # print(losses[0:10])
    losses = (losses-losses.min())/(losses.max()-losses.min())    

    # print(losses[0:10])
    # fit a two-component GMM to the loss
    input_loss = losses.reshape(-1,1)
    gmm = GaussianMixture(n_components=2,max_iter=10,tol=1e-2,reg_covar=1e-3)
    gmm.fit(input_loss)
    prob = gmm.predict_proba(input_loss) 
    prob = prob[:,gmm.means_.argmin()]

    uncertainty = (uncertainty-uncertainty.min())/(uncertainty.max()-uncertainty.min())
    input_uncertainty = uncertainty.reshape(-1, 1)
    gmm = GaussianMixture(n_components=2,max_iter=10,tol=1e-2,reg_covar=1e-3)
    gmm.fit(input_uncertainty)
    prob_reliable = gmm.predict_proba(input_uncertainty) 
    prob_reliable = prob_reliable[:,gmm.means_.argmin()]

    return prob, prob_reliable, embeddings

def linear_rampup(current, warm_up, rampup_length=16):
    current = np.clip((current-warm_up) / rampup_length, 0.0, 1.0)
    return args.lambda_u*float(current)

class SemiLoss(object):
    def __call__(self, outputs_x, targets_x, outputs_u, targets_u, epoch, warm_up):
        probs_u = torch.softmax(outputs_u, dim=1)

        Lx = -torch.mean(torch.sum(F.log_softmax(outputs_x, dim=1) * targets_x, dim=1))
        Lu = torch.mean((probs_u - targets_u)**2)

        return Lx, Lu, linear_rampup(epoch,warm_up)

class NegEntropy(object):
    def __call__(self,outputs):
        probs = torch.softmax(outputs, dim=1)
        return torch.mean(torch.sum(probs.log()*probs, dim=1))

def create_model(device):
    model = InceptionResNetV2(num_classes=args.num_class)
    model = model.to(device)
    return model

stats_log=open('./checkpoint/%s'%(args.id)+'_stats.txt','w') 
test_log=open('./checkpoint/%s'%(args.id)+'_acc.txt','w')     

warm_up=1

loader = dataloader.webvision_dataloader(batch_size=args.batch_size,num_workers=8,root_dir=args.data_path,log=stats_log, num_class=args.num_class)

print('| Building net')
net1 = create_model(device1)
net2 = create_model(device2)
cudnn.benchmark = True

criterion = SemiLoss()
optimizer1 = optim.SGD(net1.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
optimizer2 = optim.SGD(net2.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)

CE = nn.CrossEntropyLoss(reduction='none')
CEloss = nn.CrossEntropyLoss()
conf_penalty = NegEntropy()

all_loss = [[],[]] # save the history of losses from two networks
acc_meter = torchnet.meter.ClassErrorMeter(topk=[1,5], accuracy=True)

eval_loader = loader.run('eval_train')  
web_valloader = loader.run('test')
imagenet_valloader = loader.run('imagenet')
warmup_trainloader = loader.run('warmup')

print("eval loader {} examples".format(len(eval_loader.dataset)))
print("web_valloader {} examples".format(len(web_valloader.dataset)))
print("imagenet_valloader {} examples".format(len(imagenet_valloader.dataset)))

if args.is_load:
    load_dict = torch.load("checkpoint_origin/webvision/all_state_epoch={}.pt".format(args.load_epoch))
    net1.load_state_dict(load_dict['net1'])
    net2.load_state_dict(load_dict['net2'])
    optimizer1.load_state_dict(load_dict['optimizer1'])
    optimizer2.load_state_dict(load_dict['optimizer2'])
    epoch = load_dict['epoch']
    del load_dict
    web_acc = test(epoch, net1, net2, device1, device2, web_valloader)
    imagenet_acc = test(epoch, net1, net2, device1, device2, imagenet_valloader)
    print("After loading weight:\n| Test Epoch #%d\t WebVision Acc: %.2f%% (%.2f%%) \t ImageNet Acc: %.2f%% (%.2f%%)\n" % (
    epoch, web_acc[0], web_acc[1], imagenet_acc[0], imagenet_acc[1]))
    test_log.write('Epoch:%d \t WebVision Acc: %.2f%% (%.2f%%) \t ImageNet Acc: %.2f%% (%.2f%%)\n' % (
    epoch, web_acc[0], web_acc[1], imagenet_acc[0], imagenet_acc[1]))
    test_log.flush()
    print('\n==== net 1 evaluate training data loss ====')
    prob1, prob_reliable1, embeddings1 = eval_train(eval_loader, net1, device1, 'net1')
    print('\n==== net 2 evaluate training data loss ====')
    prob2, prob_reliable2, embeddings2 = eval_train(eval_loader, net2, device2, 'net2')

if args.amp:
    net1, optimizer1 = amp.initialize(
        net1, optimizer1, opt_level="O1")
    net2, optimizer2 = amp.initialize(
        net2, optimizer2, opt_level="O1")
for epoch in range(args.num_epochs+1):
    if epoch <= args.load_epoch:
        continue
    lr=args.lr
    if epoch >= 40:
        lr /= 10      
    for param_group in optimizer1.param_groups:
        param_group['lr'] = lr       
    for param_group in optimizer2.param_groups:
        param_group['lr'] = lr              
    
    if epoch<warm_up:
        print('Warmup Net1')
        warmup(epoch,net1,optimizer1,warmup_trainloader, device1)
        print('\nWarmup Net2')
        warmup(epoch,net2,optimizer2,warmup_trainloader, device2)
   
    else:                
        pred1 = (prob1 > args.p_threshold)      
        pred2 = (prob2 > args.p_threshold)      

        pred_reliable1 = (prob_reliable1 > args.p_threshold)      
        pred_reliable2 = (prob_reliable2 > args.p_threshold)   

        if epoch == warm_up or epoch == args.load_epoch + 1:
            bpl_trainloader1 = loader.run('bpl', pred2, prob2, embeddings2, pred_reliable2, prob_reliable2, knn_device=device1)
            bpl_trainloader2 = loader.run('bpl', pred1, prob1, embeddings1, pred_reliable1, prob_reliable1, knn_device=device2)
        else:
            bpl_trainloader1.dataset.create_bpl(
                pred2, prob2, embeddings2,
                pred_reliable2, prob_reliable2,
                knn_device=device2,
                log=loader.log)
            bpl_trainloader1.dataset.create_bpl(
                pred1, prob1, embeddings1,
                pred_reliable1, prob_reliable1,
                knn_device=device1,
                log=loader.log
            )
        with TimeUse(name="Train BPL", measure="m"):
            train_bpl(epoch,net1, copy.deepcopy(net2).to(device1), optimizer1, bpl_trainloader1, device1, 'net1')
            train_bpl(epoch,net2, copy.deepcopy(net1).to(device2), optimizer2, bpl_trainloader2, device2, 'net2')
    
    with TimeUse("Test Webvision and ImageNet", measure="m"):
        web_acc = test(epoch, net1, net2, device1, device2, web_valloader)
        imagenet_acc = test(epoch, net1, net2, device1, device2, imagenet_valloader)
    print("\n| Test Epoch #%d\t WebVision Acc: %.2f%% (%.2f%%) \t ImageNet Acc: %.2f%% (%.2f%%)\n"%(epoch,web_acc[0],web_acc[1],imagenet_acc[0],imagenet_acc[1]))  
    test_log.write('Epoch:%d \t WebVision Acc: %.2f%% (%.2f%%) \t ImageNet Acc: %.2f%% (%.2f%%)\n'%(epoch,web_acc[0],web_acc[1],imagenet_acc[0],imagenet_acc[1]))
    test_log.flush()  

    with TimeUse("Evaluation Train Data", measure="m"):
        print('\n==== net 1 evaluate training data loss ====')
        prob1, prob_reliable1, embeddings1 = eval_train(eval_loader, net1, device1, 'net1')
        print('\n==== net 2 evaluate training data loss ====')
        prob2, prob_reliable2, embeddings2 = eval_train(eval_loader, net2, device2, 'net2')
    # torch.save(all_loss,'./checkpoint/%s.pth.tar'%(args.id))
    if epoch >= args.save_epoch and epoch % args.save_epoch_interval == 0:
        save_dict = {
            "epoch": epoch,
            "net1": net1.state_dict(),
            "net2": net2.state_dict(),
            "optimizer1": optimizer1.state_dict(),
            "optimizer2": optimizer2.state_dict()
        }
        save_path = "checkpoint_pl/webvision/"
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        torch.save(save_dict, save_path + "all_state_epoch={}.pt".format(epoch))
